#!/usr/bin/env python3
"""
WebSocket测试客户端 - 通过代理连接WebSocket服务
"""
import socket
import base64
import hashlib
import threading
import time

class SimpleWebSocketClient:
    def __init__(self, proxy_host='127.0.0.1', proxy_port=8080):
        self.proxy_host = proxy_host
        self.proxy_port = proxy_port
        
    def create_websocket_key(self):
        """生成WebSocket密钥"""
        import os
        key = base64.b64encode(os.urandom(16)).decode('utf-8')
        return key
    
    def connect_to_websocket(self, ws_url, ws_host, ws_port=80):
        """通过代理连接到WebSocket服务器"""
        try:
            print(f"🔌 连接到代理服务器 {self.proxy_host}:{self.proxy_port}")
            
            # 连接到代理服务器
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.connect((self.proxy_host, self.proxy_port))
            
            # 生成WebSocket升级请求
            ws_key = self.create_websocket_key()
            
            upgrade_request = f"""GET {ws_url} HTTP/1.1\r
Host: {ws_host}\r
Upgrade: websocket\r
Connection: Upgrade\r
Sec-WebSocket-Key: {ws_key}\r
Sec-WebSocket-Version: 13\r
User-Agent: WebSocket-Test-Client/1.0\r
\r
"""
            
            print(f"📤 发送WebSocket升级请求到 {ws_host}")
            print("升级请求内容:")
            print(upgrade_request)
            
            # 发送升级请求
            sock.send(upgrade_request.encode('utf-8'))
            
            # 接收响应
            response = sock.recv(4096).decode('utf-8')
            print("📥 收到响应:")
            print(response)
            
            # 检查升级是否成功
            if '101 Switching Protocols' in response:
                print("✅ WebSocket升级成功!")
                return sock
            else:
                print("❌ WebSocket升级失败")
                sock.close()
                return None
                
        except Exception as e:
            print(f"❌ 连接错误: {e}")
            return None
    
    def test_echo_websocket(self):
        """测试连接到公共WebSocket echo服务"""
        print("🧪 测试连接到WebSocket echo服务")
        print("=" * 50)
        
        # 使用公共的WebSocket测试服务
        ws_host = "echo.websocket.org"
        ws_url = "/"
        ws_port = 80
        
        sock = self.connect_to_websocket(ws_url, ws_host, ws_port)
        
        if sock:
            print("🎉 WebSocket连接建立成功!")
            print("🔍 检查代理服务器日志以查看WebSocket请求处理过程")
            
            # 发送一个简单的文本消息 (WebSocket frame format)
            message = "Hello WebSocket through Proxy!"
            
            # 简单的WebSocket文本帧格式 (客户端需要mask)
            # 这里为了简化测试，不发送实际的WebSocket数据帧
            print(f"💬 模拟发送消息: {message}")
            
            # 关闭连接
            sock.close()
            print("🔚 WebSocket连接已关闭")
        else:
            print("❌ WebSocket连接失败")

def main():
    print("🔌 WebSocket代理测试客户端")
    print("=" * 60)
    print()
    
    client = SimpleWebSocketClient()
    
    print("📝 说明:")
    print("   此测试客户端会通过IDEA代理连接WebSocket服务")
    print("   主要用于验证代理服务器的WebSocket处理能力")
    print("   请确保IDEA代理服务器正在运行 (idea_proxy.py)")
    print()
    
    # 等待用户确认
    input("按回车键开始测试...")
    print()
    
    try:
        client.test_echo_websocket()
    except KeyboardInterrupt:
        print("\n🛑 测试被用户中断")
    
    print()
    print("🎯 测试完成! 请检查:")
    print("   1. 代理服务器控制台输出")
    print("   2. Dashboard中的WebSocket请求显示")
    print("   3. 日志文件中的WebSocket相关记录")

if __name__ == "__main__":
    main()